

import pyecharts.options as opts
from pyecharts.charts import Graph
from pyecharts.render import make_snapshot

# 绘图类对象
from snapshot_phantomjs import snapshot


class PGraph:
    def __init__(self, var_nodes=None, var_links=None):
        if var_links is None:
            var_links = []
        if var_nodes is None:
            var_nodes = []
        self.graph = None
        self.links = None
        self.nodes = None
        self.set_nodes(var_nodes)
        self.set_links(var_links)
        self.title = 'New Graph'
        self.img_name = 'graph.png'

    # 定义节点
    def set_nodes(self, var_nodes=None):
        if var_nodes is None:
            var_nodes = []
        self.nodes = var_nodes

    # 定义边
    def set_links(self, node_links=None):
        if node_links is None:
            node_links = []
        self.links = node_links

    # 初始化点，创建点实例
    def __init_nodes__(self):
        nodes_tmp = self.nodes
        nodes_real = []
        if len(nodes_tmp) is 0:
            # nodes = [
            #     opts.GraphNode(name="v1", symbol_size=20),
            #     opts.GraphNode(name="v2", symbol_size=20),
            #     opts.GraphNode(name="v3", symbol_size=20),
            #     opts.GraphNode(name="v4", symbol_size=20),
            #     opts.GraphNode(name="v5", symbol_size=20),
            # ]
            nodes_real = [
                opts.GraphNode(name="无数据", symbol_size=20)
            ]

        else:
            for node in nodes_tmp:
                nodes_real.append(opts.GraphNode(name=node, symbol_size=10))

        return nodes_real

    # 初始化边，创建边实例
    def __init_links__(self):

        links_real = []
        links_tmp = self.links

        link_num = len(links_tmp)
        if link_num is not 0:
            for i in range(link_num):
                links_real.append(opts.GraphLink(source=links_tmp[i][0],
                                                 target=links_tmp[i][1],
                                                 value=links_tmp[i][2],
                                                 linestyle_opts=opts.LineStyleOpts(color='black')))
        return links_real

    # 渲染
    def render(self, title=None, name=None, node_objects=None, link_objects=None, repulsion=None):
        print(
            f"-->title：{title if title is not None else 'default'}\n-->filename:{name if name is not None else 'default'}\n-->repulsion:{repulsion if repulsion is not None else 'default'}")
        self.graph = Graph()
        nodes_real = node_objects if node_objects is not None else self.__init_nodes__()
        links_real = link_objects if link_objects is not None else self.__init_links__()
        self.graph.add(
            "",
            nodes_real,
            links_real,
            repulsion=4000 if repulsion is None else repulsion,
            edge_label=opts.LabelOpts(
                is_show=True, position="middle", formatter='{c}'
            )
        )

        if title is not None:
            self.title = title
        if name is not None:
            self.img_name = name
        self.graph.set_global_opts(title_opts=opts.TitleOpts(title=self.title))
        make_snapshot(snapshot, self.graph.render(), output_name=self.img_name, delay=2)

    # 标记特殊路径并重新渲染
    # route接收的是要标记路径的索引列表，color接收的是字符串
    def mark_route_and_render(self, route_index, color='red', title=None, repulsion=None, name=None):
        print(f"-->Mark_color：{color if color is not None else 'default'}")
        routes = route_index
        links_tmp = self.links
        links_real = []
        for i in range(len(links_tmp)):
            if i in routes:
                links_real.append(opts.GraphLink(source=links_tmp[i][0],
                                                 target=links_tmp[i][1],
                                                 value=links_tmp[i][2],
                                                 linestyle_opts=opts.LineStyleOpts(color=color)))
            else:
                links_real.append(opts.GraphLink(source=links_tmp[i][0],
                                                 target=links_tmp[i][1],
                                                 value=links_tmp[i][2],
                                                 linestyle_opts=opts.LineStyleOpts(color='black')))

        self.title = title if title is not None else self.title
        self.render(title=title, link_objects=links_real, repulsion=repulsion, name=name)


if __name__ == '__main__':
    pGraph = PGraph()
    nodes = [
        'v1',
        'v2',
        'v3',
        'v4',
        'v5'
    ]

    links = [
        ['v1', 'v2', 5],
        ['v1', 'v3', 1],
        ['v3', 'v4', 3],
        ['v2', 'v4', 4],
        ['v4', 'v5', 10],
        ['v3', 'v5', 12]
    ]

    pGraph.set_nodes(var_nodes=nodes)
    pGraph.set_links(node_links=links)
    pGraph.render()
    x = input('选择最短路径')
    if x is '1':
        pGraph.mark_route_and_render(route_index=[
            0,
            3,
            4
        ])
